import sqlite3
import hydmaker
import os
import itertools
import traceback, code
import re
import csv
import copy
import threading

import wx
import matplotlib
matplotlib.use('WXAgg')
from matplotlib.figure import Figure
from matplotlib.backends.backend_wxagg import \
        FigureCanvasWxAgg as FigCanvas, \
        NavigationToolbar2WxAgg as NavigationToolbar

class MasterFrame(wx.Frame):

    title = 'Hydmaker volume correction'

    def __init__(self, dbpath, window_size, basevol_scale, vol_rating_weight, side, area_based_rating):
        wx.Frame.__init__(self, None, -1, self.title)
        self.dbpath = dbpath
        self.cnxn = sqlite3.connect(dbpath)
        self.cnxn.isolation_level = None
        self.cursor = self.cnxn.cursor()
        self.cursor.execute("PRAGMA synchronous=OFF") # don't wait for disk writes to complete before continuing
        self.cursor.execute("PRAGMA count_changes=OFF") # don't count num. rows affected by DELETE, INSERT, or UPDATE
        self.cursor.execute("PRAGMA journal_mode=OFF") # turn off journaling
        self.Bind(wx.EVT_CLOSE, self.OnClose)

        self.cursor.execute("SELECT timestep_secs FROM settings")
        self.timestep_secs = self.cursor.fetchone()[0]
        
        self.default_settings = {
                'window_size' : window_size,
                'basevol_scale' : basevol_scale,
                'vol_rating_weight' : vol_rating_weight,
                'side' : side,
                'area_based_rating' : area_based_rating
        }

        self.cursor.execute("SELECT max(timestep) FROM volumes2_b")
        ntimesteps = self.cursor.fetchone()[0]

        self.cursor.execute("PRAGMA table_info(correction_settings)")
        if not self.cursor.fetchall():
            self.cursor.execute("""
                CREATE TABLE correction_settings (
                    swmm integer,
                    window_size integer,
                    basevol_scale real,
                    vol_rating_weight real,
                    side integer,
                    area_based_rating integer
                )
            """)

        self.cursor.execute("DELETE FROM correction_settings")

        self.panel = wx.Panel(self)
        self.dpi = 100
        self.fig = Figure(figsize=(15, 5), dpi=self.dpi)
        self.canvas = FigCanvas(self.panel, -1, self.fig)

        self.menubar = wx.MenuBar()
        
        menu_file = wx.Menu()
        m_load = menu_file.Append(-1, "&Load correction settings")
        self.Bind(wx.EVT_MENU, self.load_correction_settings, m_load)
        m_export = menu_file.Append(-1, "&Export correction settings")
        self.Bind(wx.EVT_MENU, self.export_correction_settings, m_export)
        m_export_hyd = menu_file.Append(-1, "&Update database and export .hyd file")
        self.Bind(wx.EVT_MENU, self.export_hyd, m_export_hyd)
        menu_file.AppendSeparator()
        m_exit = menu_file.Append(-1, "E&xit")
        self.Bind(wx.EVT_MENU, self.OnClose, m_exit)
        
        self.menubar.Append(menu_file, "&File")
        self.SetMenuBar(self.menubar)

        self.statusbar = self.CreateStatusBar()

        self.fig.subplots_adjust(top=.9, bottom=.1, right=.98, left=.05)

        gs = matplotlib.gridspec.GridSpec(2, 1, height_ratios=[3,1])
        self.axes = self.fig.add_subplot(gs[0])
        self.bv_axes = self.fig.add_subplot(gs[1], sharex=self.axes)

        self.cursor.execute("SELECT DISTINCT swmm FROM volumes2_b WHERE swmm IS NOT NULL")
        self.swmm_segs = [row[0] for row in self.cursor.fetchall()]

        if not self.swmm_segs:
            raise Exception("That database is empty.")

        self.current_seg = None

        seg_choice_lbl = wx.StaticText(self.panel, -1, "SWMM Segment")
        self.seg_choice = wx.Choice(self.panel, -1, (100, 50), choices=[str(x) for x in self.swmm_segs])
        self.panel.Bind(wx.EVT_CHOICE, self.on_seg_choice, self.seg_choice)

        self.prev_seg_btn = wx.Button(self.panel, -1, 'Previous')
        self.prev_seg_btn.Disable()
        self.Bind(wx.EVT_BUTTON, self.on_prev, self.prev_seg_btn)
        self.next_seg_btn = wx.Button(self.panel, -1, 'Next')
        self.Bind(wx.EVT_BUTTON, self.on_next, self.next_seg_btn)

        class WinsizeValidator(wx.PyValidator):
            def __init__(self, max_size):
                wx.PyValidator.__init__(self)
                self.max_size = max_size

            def Clone(self):
                return WinsizeValidator(self.max_size)

            def Validate(self, win):
                textCtrl = self.GetWindow()
                msg = "The neighborhood width needs to be greater than 2 and less than " + str(self.max_size)
                try:
                    value = int(textCtrl.GetValue())
                except:
                    wx.MessageBox(msg)
                    return False

                if value < 2 or value > self.max_size:
                    wx.MessageBox(msg)
                    return False
                else:
                    textCtrl.SetValue(int(value))
                    return True
            
            def TransferToWindow(self):
                return True

            def TransferFromWindow(self):
                return True

        class BVScaleValidator(wx.PyValidator):
            def __init__(self):
                wx.PyValidator.__init__(self)

            def Clone(self):
                return BVScaleValidator()

            def Validate(self, win):
                textCtrl = self.GetWindow()
                msg = "DWV tolerance rating has to be a number greater than zero."
                try:
                    value = float(textCtrl.GetValue())
                except:
                    wx.MessageBox(msg)
                    return False


                if value < 0:
                    wx.MessageBox(msg)
                    return False
                else:
                    return True
            
            def TransferToWindow(self):
                return True

            def TransferFromWindow(self):
                return True

        class VolRatingWeightValidator(wx.PyValidator):
            def __init__(self):
                wx.PyValidator.__init__(self)

            def Clone(self):
                return VolRatingWeightValidator()

            def Validate(self, win):
                textCtrl = self.GetWindow()
                msg = "QinQout rating bias should be a number between 0 and 1"
                try:
                    value = float(textCtrl.GetValue())
                except:
                    wx.MessageBox(msg)
                    return False

                if value < 0 or value > 1:
                    wx.MessageBox(msg)
                    return False
                else:
                    return True
            
            def TransferToWindow(self):
                return True

            def TransferFromWindow(self):
                return True

        winsize_text_lbl = wx.StaticText(self.panel, -1, 'DWV rating neighborhood size:')
        self.winsize_text = wx.TextCtrl(self.panel, -1, "", size=(40, -1), validator=WinsizeValidator(int(ntimesteps/2)))
        self.Bind(wx.EVT_TEXT, self.on_setting_change, self.winsize_text)

        side_text_lbl = wx.StaticText(self.panel, -1, 'DWV rating neighborhood position:')
        choices = ['left', 'center', 'right']
        self.side_rb = wx.RadioBox(self.panel, -1, '', wx.DefaultPosition, wx.DefaultSize, choices, 3, wx.RA_SPECIFY_COLS | wx.NO_BORDER)
        self.Bind(wx.EVT_RADIOBOX, self.on_setting_change, self.side_rb)

        basevol_scale_lbl = wx.StaticText(self.panel, -1, 'DWV tolerance volume:')
        self.basevol_scale_text = wx.TextCtrl(self.panel, -1, "", size=(40, -1), validator=BVScaleValidator())
        self.Bind(wx.EVT_TEXT, self.on_setting_change, self.basevol_scale_text)

        vol_rating_weight_lbl = wx.StaticText(self.panel, -1, 'QinQout rating bias:')
        self.vol_rating_weight_text = wx.TextCtrl(self.panel, -1, "", size=(40, -1), validator=VolRatingWeightValidator())
        self.Bind(wx.EVT_TEXT, self.on_setting_change, self.vol_rating_weight_text)
        self.vol_rating_weight_text.Disable()

        self.area_based_rating_cb = wx.CheckBox(self.panel, -1, 'Use area based rating')
        def on_area_based_rating_cb(event):
            if event.Checked():
                self.vol_rating_weight_text.Disable()
            else:
                self.vol_rating_weight_text.Enable()
            self.on_setting_change(event)

        self.Bind(wx.EVT_CHECKBOX, on_area_based_rating_cb, self.area_based_rating_cb)
    
        self.update_btn = wx.Button(self.panel, -1, "Update")
        self.Bind(wx.EVT_BUTTON, self.on_update_button, self.update_btn)
        self.update_btn.Disable()

        self.toolbar = NavigationToolbar(self.canvas)

        self.vbox = wx.BoxSizer(wx.VERTICAL)
        self.vbox.Add(self.canvas, 1, wx.LEFT | wx.TOP | wx.GROW)
        self.vbox.Add(self.toolbar, 0, wx.EXPAND)
        self.vbox.AddSpacer(10)

        self.hbox = wx.BoxSizer(wx.HORIZONTAL)
        flags = wx.ALIGN_LEFT | wx.ALL | wx.ALIGN_CENTER_VERTICAL

        self.hbox.Add(seg_choice_lbl, 0, border=3, flag=flags)
        self.hbox.Add(self.seg_choice, 0, border=3, flag=flags)
        self.hbox.Add(self.prev_seg_btn, border=3, flag=flags)
        self.hbox.Add(self.next_seg_btn, border=3, flag=flags)
        self.hbox.Add(self.update_btn, 0, border=3, flag=flags)

        self.hbox.Add(winsize_text_lbl, 0, border=3, flag=flags)
        self.hbox.Add(self.winsize_text, 0, border=3, flag=flags)

        self.hbox.Add(side_text_lbl, 0, border=3, flag=flags)
        self.hbox.Add(self.side_rb, 0, border=3, flag=flags)

        self.hbox.Add(basevol_scale_lbl, 0, border=3, flag=flags)
        self.hbox.Add(self.basevol_scale_text, 0, border=3, flag=flags)

        self.hbox.Add(self.area_based_rating_cb, 0, border=3, flag=flags)
        self.hbox.Add(vol_rating_weight_lbl, 0, border=3, flag=flags)
        self.hbox.Add(self.vol_rating_weight_text, 0, border=3, flag=flags)


        self.vbox.Add(self.hbox, 0, flag=wx.ALIGN_LEFT | wx.TOP)

        self.plot_changed = False

        self.panel.SetSizer(self.vbox)
        self.vbox.Fit(self)
    
    def update_next_prev_btns(self):
        current_seg_selection = self.seg_choice.GetSelection()
        if current_seg_selection > 0:
            self.prev_seg_btn.Enable()
        else:
            self.prev_seg_btn.Disable()

        if current_seg_selection < len(self.swmm_segs) - 1:
            self.next_seg_btn.Enable()
        else:
            self.next_seg_btn.Disable()

    def on_next(self, event):
        current_seg_selection = self.seg_choice.GetSelection()
        if current_seg_selection < len(self.swmm_segs) - 1:
            self.seg_choice.SetSelection(current_seg_selection + 1)
            self.update_next_prev_btns()
            self.load_seg(self.swmm_segs[current_seg_selection + 1])


    def on_prev(self, event):
        current_seg_selection = self.seg_choice.GetSelection()
        if current_seg_selection > 0:
            self.seg_choice.SetSelection(current_seg_selection - 1)
            self.update_next_prev_btns()
            self.load_seg(self.swmm_segs[current_seg_selection - 1])

    def on_setting_change(self, event):
        self.update_btn.Enable()

    def on_seg_choice(self, event):
        seg_choice = int(event.GetString())
        seg_choice_idx = self.swmm_segs.index(seg_choice)
        self.update_next_prev_btns()

        self.load_seg(seg_choice)

    def get_vol_and_qvl(self, seg, cursor=None):
        if not cursor:
            cursor = self.cursor

        cursor.execute("""
            SELECT vol, flow, velocity
            FROM volumes2_b
            WHERE swmm = ?
            ORDER BY timestep
        """, (seg,))
        vol, flow, vel = zip(*cursor.fetchall())

        cursor.execute("""
            SELECT DISTINCT length
            FROM conduit_params2 WHERE swmm = ?
        """, (seg,))

        l = cursor.fetchone()[0]
        qvl_vol = [f/v * l for f, v in zip(flow, vel)]

        return vol, qvl_vol

    def set_correction_settings(self, seg, cursor=None):
        if not cursor:
            cursor = self.cursor
        cursor.execute("SELECT * FROM correction_settings WHERE swmm = ?", (seg,))
        result = cursor.fetchone()
        if not result:
            cursor.execute("""
                INSERT INTO correction_settings (swmm, window_size, basevol_scale, vol_rating_weight, side, area_based_rating)
                VALUES (:swmm, :window_size, :basevol_scale, :vol_rating_weight, :side, :area_based_rating)
            """, dict(self.default_settings.items() + [('swmm', seg)]))

            self.current_settings = {key:value for key, value in self.default_settings.items()}
        else:
            self.current_settings = dict(zip([x[0] for x in cursor.description], result))

    def load_seg(self, seg_choice, refresh=False):
        self.statusbar.SetStatusText(('Refreshing' if refresh else 'Loading') + " segment " + str(seg_choice) + " ...")
        if seg_choice != self.current_seg or refresh:
            if self.current_seg is not None and not refresh:
                self.update_database_settings(self.current_seg, self.current_settings)

            self.update_btn.Disable()
            self.current_seg = seg_choice

            self.set_correction_settings(self.current_seg)
            self.vol, self.qvl_vol = self.get_vol_and_qvl(self.current_seg)
            self.update_settings_widgets()
            self.corrected_vol = self.correct(self.vol, self.qvl_vol)

            self.plot_changed = False
            if refresh:
                self.statusbar.SetStatusText("Segment " + str(seg_choice) + ' refreshed.')
            else:
                self.statusbar.SetStatusText("Showing segment " + str(seg_choice) + '.')
            self.draw_figure()

    def draw_figure(self):
        self.axes.clear()
        x = range(len(self.vol))
        self.axes.plot(x, self.vol, label='(Qin - Qout) * timestep')
        self.axes.plot(x, self.qvl_vol, label='(Q/v) * l')
        self.axes.plot(x, self.corrected_vol, label='corrected')

        self.axes.legend(bbox_to_anchor=(0., 1.02, 1., .102), loc=3, ncol=3, borderaxespad=0., prop={'size':12})
        self.axes.set_ylabel('volume (meters-cubed)')
        self.axes.grid(True)

        self.bv_axes.clear()
        self.bv_axes.plot(self.vol_ratings, label='vol_rating')
        self.bv_axes.plot(self.qvl_ratings, label='qvl_rating')
        self.bv_axes.plot(x, self.basevol_ratings, label='basevol_rating')

        self.bv_axes.set_xlabel("timestep")
        self.bv_axes.set_ylabel("DWV tol. rating")
        self.bv_axes.grid(True)

        self.canvas.draw()

    def update_settings(self):
        self.current_settings['window_size'] = int(self.winsize_text.GetValue())
        self.current_settings['side'] = self.side_rb.GetSelection() - 1
        self.current_settings['basevol_scale'] = float(self.basevol_scale_text.GetValue())
        self.current_settings['area_based_rating'] = bool(self.area_based_rating_cb.GetValue())
        self.current_settings['vol_rating_weight'] = float(self.vol_rating_weight_text.GetValue())

    def update_settings_widgets(self):
        self.winsize_text.SetValue(str(self.current_settings['window_size']))
        self.side_rb.SetSelection(self.current_settings['side'] + 1)
        self.basevol_scale_text.SetValue(str(self.current_settings['basevol_scale']))
        area_based_rating = self.current_settings['area_based_rating']
        self.area_based_rating_cb.SetValue(area_based_rating)

        self.vol_rating_weight_text.SetValue(str(self.current_settings['vol_rating_weight']))
        if area_based_rating:
            self.vol_rating_weight_text.Disable()
        else:
            self.vol_rating_weight_text.Enable()
        
    def update_database_settings(self, seg, settings):
        self.cursor.execute("""
            UPDATE correction_settings
            SET window_size = :window_size,
                basevol_scale = :basevol_scale,
                vol_rating_weight = :vol_rating_weight,
                area_based_rating = :area_based_rating,
                side = :side
            WHERE swmm = :swmm
        """, dict(settings.items() + [('swmm', seg)]))

    def on_update_button(self, event):
        self.statusbar.SetStatusText("Updating segment " + str(self.current_seg) + ".")
        self.update_settings()
        self.update_database_settings(self.current_seg, self.current_settings)
        self.corrected_vol = self.correct(self.vol, self.qvl_vol)
        self.draw_figure()
        self.update_btn.Disable()
        self.statusbar.SetStatusText("Segment " + str(self.current_seg) + " udpated.")

    def correct(self, vol, qvl_vol, return_corrected_vol=True):
        vol = list(vol)
        
        total_correction = 0
        vol_correction = [0 for _ in range(len(vol))]

        corrected_vol = copy.copy(vol)

        i = 0
        running_total = 0

        self.basevol_ratings = [0 for _ in range(len(vol))]

        self.vol_ratings = [0 for _ in range(len(vol))]
        self.qvl_ratings = [0 for _ in range(len(vol))]

        qvl_area = sum([x * self.timestep_secs for x in qvl_vol])
        vol_area = sum([x * self.timestep_secs for x in vol])

        window_size = self.current_settings['window_size']
        basevol_scale = self.current_settings['basevol_scale']
        vol_rating_weight = self.current_settings['vol_rating_weight']
        area_based_rating = self.current_settings['area_based_rating']
        side = self.current_settings['side']
        while i < len(vol):
            if side == 0:
                lowerbound = max(0, i - int(window_size/2.0))
                upperbound = min(len(vol), i + int(window_size/2.0))
            elif side == -1:
                lowerbound = max(0, i - window_size)
                upperbound = i
                if lowerbound == upperbound:
                    if i + 1 < len(vol):
                        upperbound += 1
                    else:
                        lowerbound = lowerbound - 1
            else:
                lowerbound = i
                upperbound = min(len(vol), i + window_size)
                if lowerbound == upperbound:
                    lowerbound = lowerbound - 1

            vol_diff = max(vol[lowerbound:upperbound]) - min(vol[lowerbound:upperbound])
            vol_rating = min(1, vol_diff / basevol_scale)
            self.vol_ratings[i] = vol_rating

            qvl_diff = max(qvl_vol[lowerbound:upperbound]) - min(qvl_vol[lowerbound:upperbound])
            qvl_rating = min(1, qvl_diff / basevol_scale)
            self.qvl_ratings[i] = qvl_rating

            bv_rating = vol_rating_weight * vol_rating + (1.0 - vol_rating_weight) * qvl_rating
            self.basevol_ratings[i] = bv_rating

            error = qvl_vol[i] - vol[i]
            if area_based_rating:
                if error > 0:
                    error = error - (qvl_rating if vol_area > qvl_area else vol_rating) * error
                else:
                    # consider the possibly of always using vol_rating when the error is negative
                    error = error - (vol_rating if vol_area > qvl_area else qvl_rating) * error
            else:
                error = error - bv_rating * error

            corrected_vol[i] += error
            correction = error - running_total
            vol_correction[i] = correction
            running_total += correction

            i += 1

        if return_corrected_vol:
            return corrected_vol
        else:
            return vol_correction

    def OnClose(self, event):
        self.cnxn.close()
        self.Destroy()

    def export_correction_settings(self, event):
        self.statusbar.SetStatusText("Exporting correction settings to database directory.")
        self.cursor.execute("SELECT * FROM correction_settings")
        rows = [dict(zip([x[0] for x in self.cursor.description], row)) for row in self.cursor.fetchall()]
        settings_segs = [row['swmm'] for row in rows]
        for seg in [seg for seg in self.swmm_segs if seg not in settings_segs]:
            new_row = copy.copy(self.default_settings)
            new_row['swmm'] = seg
            rows.append(new_row)

        fieldnames = ['swmm'] + self.default_settings.keys()
        out_fname = os.path.splitext(os.path.basename(self.dbpath))[0] + '_correction_settings.csv'
        outpath = os.path.join(os.path.dirname(self.dbpath), out_fname)
        rows = sorted(rows, key=lambda x: x['swmm'])
        for row in rows:
            row['area_based_rating'] = int(row['area_based_rating'])
        with open(outpath, 'w') as f:
            writer = csv.DictWriter(f, lineterminator='\n', fieldnames=fieldnames)
            writer.writeheader()
            writer.writerows(rows)

        self.statusbar.SetStatusText("Correction settings exported to database directory.")

    def export_hyd(self, event):
        self.cursor.close()
        self.cnxn.close()
        threading.Thread(target=self.finalize_and_export, args=(True,)).start()

    def finalize_and_export(self, sep_thread=False):
        if sep_thread:
            cnxn = sqlite3.connect(self.dbpath)
            cnxn.isolation_level = None
            cursor = cnxn.cursor()
            cursor.execute("PRAGMA synchronous=OFF") # don't wait for disk writes to complete before continuing
            cursor.execute("PRAGMA count_changes=OFF") # don't count num. rows affected by DELETE, INSERT, or UPDATE
            cursor.execute("PRAGMA journal_mode=OFF") # turn off journaling
        else:
            cursor = self.cursor
        self.statusbar.SetStatusText("Finalizing corrections. This may take a moment.")
        cursor.execute("SELECT timestep_secs FROM settings")
        timestep_secs = cursor.fetchone()[0]

        cursor.execute("DELETE FROM flows WHERE name LIKE '%balbucket' OR name LIKE '%nonzero'")
        try:
            cursor.execute('CREATE INDEX idx_flows2 ON flows(swmm_sink, swmm_source, timestep)')
        except:
            pass

        for seg in self.swmm_segs:
            vol, qvl_vol = self.get_vol_and_qvl(seg, cursor=cursor)
            self.set_correction_settings(seg, cursor=cursor)
            correction = self.correct(vol, qvl_vol, return_corrected_vol=False)

            balance_bucket = [c/timestep_secs for c in correction]
            update_tups = []
            cursor.execute("""
                SELECT DISTINCT wasp
                FROM segs 
                WHERE swmm = ?
            """, (seg,))
            wasp = cursor.fetchone()[0]

            update_tups = [[i, 0, seg, 0, wasp, str(seg)+'balbucket', qin] for i, qin in enumerate(balance_bucket)]

            # update the flows table with the correction inlet 
            insert_qry = """
                INSERT INTO flows (timestep, swmm_source, swmm_sink, wasp_source, wasp_sink, name, flow) 
                VALUES (?, ?, ?, ?, ?, ?, ?)
            """
            cursor.executemany(insert_qry, update_tups)

            cursor.execute("""
                SELECT timestep, SUM(flow) 
                FROM flows 
                WHERE swmm_source = 0 AND swmm_sink = ?
                GROUP BY timestep
            """, (seg,))

            total_inflows = cursor.fetchall()

            nonzero_guarantee = [0 for _ in range(len(balance_bucket))]

            for bb, tot_inf in zip(balance_bucket, total_inflows):
                if balance_bucket != 0 and total_inflows[1] == 0:
                    nonzero_guarantee[total_inflows[0]] = 0.0001

            update_tups = [[i, 0, seg, 0, wasp, str(seg) + 'nonzero', x] for i, x in enumerate(nonzero_guarantee)]
            cursor.executemany(insert_qry, update_tups)

        self.statusbar.SetStatusText("Corrections finalized, recalculating volumes.")
        print 'ok'

        # recalculate swmm volumes with the new dummy inlets
        hydmaker.calc_swmm_vols(cursor, timestep_secs)

        try:
            cursor.execute("SELECT flow_weight_sums FROM settings")
        except:
            flow_weight_sums = False
        else:
            flow_weight_sums = bool(cursor.fetchone()[0])

        self.statusbar.SetStatusText("Volumes recalculated, preparing data .hyd file.")
        outpath = os.path.join(os.path.dirname(self.dbpath), 
                os.path.splitext(os.path.basename(self.dbpath))[0] + '_corrected.hyd.txt')
        hydmaker.finalize(cursor, outpath, 5, using_swmm_volume=False, flow_weight_sums=flow_weight_sums)

        if sep_thread:
            cursor.close()
            cnxn.close()

        if sep_thread:
            wx.CallAfter(self.export_hyd_complete)

    def export_hyd_complete(self):
        self.statusbar.SetStatusText("Export complete.")
        self.cnxn = sqlite3.connect(self.dbpath)
        self.cursor = self.cnxn.cursor()
        self.cursor.execute("PRAGMA synchronous=OFF") # don't wait for disk writes to complete before continuing
        self.cursor.execute("PRAGMA count_changes=OFF") # don't count num. rows affected by DELETE, INSERT, or UPDATE
        self.cursor.execute("PRAGMA journal_mode=OFF") # turn off journaling

    def load_correction_settings(self, event):
        wildcard = 'CSV file (*.csv)|*.csv'
        dlg = wx.FileDialog(self, message='Load correction settings', 
                defaultDir=os.path.dirname(self.dbpath),
                wildcard=wildcard,
                style=wx.OPEN)
        dlg.CenterOnParent()

        if dlg.ShowModal() == wx.ID_OK:
            self.statusbar.SetStatusText("Loading correction settings.")
            path = str(dlg.GetPaths()[0])

            with open(path, 'r') as f:
                fieldnames = f.readline()
                fieldnames = [n.strip() for n in fieldnames.split(',')]
                reader = csv.DictReader(f, fieldnames=fieldnames)

                self.cursor.execute("DELETE FROM correction_settings")
                rows = list(reader)
                self.cursor.executemany("""
                    INSERT INTO correction_settings (swmm, window_size, basevol_scale, vol_rating_weight, side, area_based_rating)
                    VALUES (:swmm, :window_size, :basevol_scale, :vol_rating_weight, :side, :area_based_rating)
                """, rows)

            self.statusbar.SetStatusText("Correction settings loaded.")
            if self.current_seg is not None:
                self.load_seg(self.current_seg, refresh=True)

class CorrectionParameter(object):
    def __init__(self, variable, label, value, msg, prompt, conversion, validator=None):
        self.variable = variable
        self.label = label
        self.value = value
        self.msg = msg
        self.prompt = prompt 
        self.conversion = conversion
        self.validator = validator

def prompt_correction_parameters(default=False):
    # correction_params is a dictionary of the default correction parameters.
    # the following prompts allow the user to change these parameters, if they desire the correction 
    # routine to run on their data

    correction_params = [
        CorrectionParameter(variable='window_size',
                            label='Dry weather volume rating neighborhood size',
                            value=800,  # old 600
                            msg="Sets the size of the neighborhood around a point that determines its dry weather volume rating.",
                            prompt="DWV rating neighborhood size",
                            conversion=int,
                            validator=lambda x: x > 1),
        CorrectionParameter(variable='basevol_scale',
                            label='Dry weather volume tolerance volume',
                            value=90, # old 30
                            msg="Sets the scale on which the margnitude of the range of a dry weather volume neighborhood is graded",
                            prompt='DWV tolerance volume',
                            conversion=float,
                            validator=lambda x: x > 0),
        CorrectionParameter(variable='vol_rating_weight',
                            label='QinQout rating bias',
                            value=.5,
                            msg=("When the rating is calculated as an average of the QinQout and QVL volumes, by what "
                                 "percentage should the QinQout rating be preferred? This is only used if the area based "
                                 "rating option is False."),
                            prompt='QinQout rating bias',
                            conversion=float,
                            validator=lambda x: x >= 0 and x <= 1),
        CorrectionParameter(variable='side',
                            label='Dry weather volume rating neighborhood position',
                            value=1,
                            msg="How should the neighborhood be positioned around the point?",
                            prompt='DWV rating neighborhood position (-1: left, 0: center, 1: right)',
                            conversion=int,
                            validator=lambda x: x in (-1, 0, 1)),
        CorrectionParameter(variable='area_based_rating',
                            label='Use area based rating',
                            value=False,
                            msg='Use the QinQout rating or the QVL rating depending on which curve has the larger total volume.',
                            prompt='Use area based rating (y/n)',
                            conversion=lambda x: bool(re.match('y', x)),
                            validator= lambda x: True)
        ]
    if not default:
        while True:
            msg = "There are " + str(len(correction_params)) + " parameters that you can set to change the behavior " + \
                "of the volume correction procedure."
            print(msg)
            menu = [str(i) + ' - ' + param.label + ' (current value: ' + str(param.value) + ')' 
                    for i, param in enumerate(correction_params)]
            print('\n'.join(menu))
            print("Enter the parameter number you want to change. Leave blank and press enter to continue.")
            choice = raw_input(">> ")

            if choice == '':
                break
            else:
                try:
                    choice = int(choice)
                    assert choice >= 0 and choice < len(correction_params)
                except:
                    print("Invalid choice.")
                else:
                    while True:
                        param = correction_params[choice]
                        print("Volume correction parameter " + str(choice) + ": " + param.label)

                        print("Description: " + param.msg)
                        print("Current value: " + str(param.value))
                        value = raw_input(param.prompt + ' >> ')

                        try:
                            value = param.conversion(value)
                            if param.validator:
                                assert param.validator(value)
                        except:
                            print("Invalid value.")
                        else:
                            print("Correction parameter value set.")
                            correction_params[choice].value = value
                            break

    return {param.variable : param.value for param in correction_params}


def run():
    while True:
        dbpath = raw_input('Hydmaker database path >> ')
        if os.path.exists(dbpath):
            break
        print("No such database found.")

    # by default, assume that the user does not want to correct volume
    correction_param_args = prompt_correction_parameters(default=True)
    # determine if user wants to correct the volume
    while True:
        msg = "There are several parameters that you can change to control the behavior of the volume correction " \
            + "procedure. Would you like to change them from their default values?"
        print(msg)
        try:
            reject_defaults = raw_input("Change default volume correction parameters? (y/n) >> ")[0]
            reject_defaults = bool(re.match('y', reject_defaults, re.IGNORECASE))
        except:
            print("Invalid input.")
        else:
            if reject_defaults:
                correction_param_args = prompt_correction_parameters()
            break

    app = wx.PySimpleApp()
    app.frame = MasterFrame(dbpath, **correction_param_args)
    app.frame.Show()
    app.MainLoop()

